"""Solver for the Burger's equation"""

import solver

from os.path import join
import numpy

class BurgersSolver(solver.Solver1D):
    pass

class BurgersWithSource1D(BurgersSolver):
    def __init__(self, tf=1.0, pfreq=10, sfunc=None, grid=None):
        BurgersSolver.__init__(self, tf=tf, pfreq=pfreq, nvar=1, grid=grid)

        # save the source function
        self.sfunc = sfunc

    def update(self, dt):
        # update the solution using the convection algorithm
        BurgersSolver.update(self, dt)

        # now solve the ODE
        grid = self.grid
        sfunc = self.sfunc

        if sfunc is None:
            raise NotImplementedError("Source function not defined")

        for cell in grid.cells[grid.nb:-grid.nb]:
            x = grid.xc[ cell.index ]
            q = cell.q

            # get the forcing function RHS
            sval = sfunc(x, q)

            # update the cell value using a simple Euler integration
            q[:] = q + sval * dt
